Write a custom CUDA kernel to optimize `Balanced Softmax Loss`.

Formula: Loss = -log( exp(z_target + log(n_target)) / Sum( exp(z_j + log(n_j)) ) )
Where `z` are logits, and `n_j` is the sample frequency of class `j`.
This is mathematically equivalent to adding a bias term `log(n_j)` to the logits before standard Cross Entropy.

Problem Analysis:
1. Memory Overhead: The standard implementation `F.cross_entropy(logits + log_freqs.log(), targets)` creates a full-sized `(N, C)` intermediate tensor for the adjusted logits.
2. Bandwidth Waste: This intermediate tensor is written to global memory and immediately read back.

Optimization Strategy: Fused Bias-Add Softmax-NLL Kernel

1. Pre-computation: Calculate `log_freqs = log(samples_per_class)` on the CPU/GPU once and pass it as a 1D tensor `(C,)` to the kernel.

2. One-Block-per-Row: Assign one CUDA block to process one sample (row) of the logits.

3. Shared Memory & Vectorization:
   - Use `float4` vectorization to load `logits` (from `N x C`) and `log_freqs` (from `1 x C`).
   - Since `log_freqs` is broadcasted across the batch, accessing it follows a consistent pattern (though not strictly coalesced across blocks, it hits the L2 cache efficiently).

4. Fused Logic:
   - Pass 1 (Max): Compute `val = logit[i] + log_freqs[i]`. Find Max `M` of these adjusted values.
   - Pass 2 (Sum): Compute SumExp `S = sum(exp(val - M))`.
   - Target Adjustment: Identify `z_target` and `n_target`.

5. Final Calculation:
   - `Loss = log(S) + M - (z_target + log(n_target))`
   - Write the single scalar loss per sample.
  
Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is:   
  
```python
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

BATCH_SIZE = 2048
NUM_CLASSES = 4096
SHAPE = (BATCH_SIZE, NUM_CLASSES)

REDUCTION = 'none'

class BalancedSoftmaxLoss(nn.Module):
    """
    Balanced Softmax Loss (NeurIPS 2020)
    L = -log( (n_y * e^{z_y}) / sum(n_j * e^{z_j}) )
      = CrossEntropy(logits + log(frequencies), targets)
    """
    def __init__(self, samples_per_class, reduction='mean'):
        super(BalancedSoftmaxLoss, self).__init__()
        self.reduction = reduction
        
        # 预计算 log
        samples_per_class = torch.as_tensor(samples_per_class).float()
        log_freqs = torch.log(samples_per_class)
        
        self.register_buffer('log_freqs', log_freqs)

    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        # logits: (N, C)
        # log_freqs: (C)
        
        # Adjust Logits (Broadcast add)
        adjusted_logits = logits + self.log_freqs.unsqueeze(0)
        
        loss = F.cross_entropy(adjusted_logits, targets, reduction='none')
        
        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        return loss

class Model(nn.Module):
    def __init__(self, samples_per_class, reduction='none'):
        super(Model, self).__init__()
        self.loss_fn = BalancedSoftmaxLoss(samples_per_class, reduction=reduction)
    
    def forward(self, logits, targets):
        return self.loss_fn(logits, targets)

def get_inputs():
    logits = torch.randn(SHAPE, dtype=torch.float32)
    targets = torch.randint(0, NUM_CLASSES, (BATCH_SIZE,), dtype=torch.long)
    return [logits.contiguous(), targets.contiguous()]

def get_init_inputs():
    # 模拟长尾分布样本数
    idx = torch.arange(NUM_CLASSES, dtype=torch.float32)
    samples_per_class = 5000.0 * (0.1 ** (idx / (NUM_CLASSES - 1)))
    samples_per_class = torch.maximum(samples_per_class, torch.tensor(1.0))
    
    return [samples_per_class, REDUCTION]